import os
import torch

from torchvision.io import read_image
from torchvision.ops.boxes import masks_to_boxes
from torchvision import tv_tensors
from torchvision.transforms.v2 import functional as F
from torch.utils.data import Dataset
from data_convert import get_transform

class PennFudanDataset(Dataset):
    def __init__(self,root,transforms):
        self.root = root
        self.transforms = transforms
        self.imgs = list(sorted(os.listdir(os.path.join(root,"PNGImages"))))
        self.masks = list(sorted(os.listdir(os.path.join(root,"PedMasks"))))

    def __getitem__(self,idx):
        #load image and mask
        img_path = os.path.join(self.root,"PNGImages",self.imgs[idx])
        mask_path = os.path.join(self.root,"PedMasks",self.masks[idx])
        img = read_image(img_path)
        mask = read_image(mask_path)
        obj_ids = torch.unique(mask)

        # first id is the background, so remove it
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)
        # print(obj_ids)
        # print(obj_ids[:, None, None])
        # print(mask.shape)
        
        # split the color-encoded mask into a set of binary masks
        masks = (mask == obj_ids[:,None,None]).to(dtype=torch.uint8)
        print(f"masks:{masks.shape}")

        # get bounding box coordinates for each mask
        boxes = masks_to_boxes(masks)
        # print(f"boxes:{boxes}")

        # there is only one class
        labels = torch.ones((num_objs,), dtype=torch.int64)

        image_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # print(f"point:{boxes[:, 3]}")

        # suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # Wrap sample and targets into torchvision tv_tensors:
        img = tv_tensors.Image(img)
        target = {}
        target["boxes"] = tv_tensors.BoundingBoxes(boxes, format="XYXY", canvas_size=F.get_size(img))
        target["masks"] = tv_tensors.Mask(masks)
        temp = target["masks"]
        print(f"target masks:{temp.shape}")
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms is not None:
            img, target = self.transforms(img, target)

        return img,target
    
    def __len__(self):
        return len(self.imgs)

def main():
    pendata = PennFudanDataset('data/PennFudanPed',get_transform(False))
    pendata[0]

if __name__ == '__main__':
    main()